#!/usr/bin/python3
import os
import cv2
import numpy as np
import struct

face_select_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/SelectMerge/MRFeature.select.txt'
video_label_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/video_label_list.txt'
file_list_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/feats_list.base.txt'
feature_file_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/feats.base.dat'

feat_dim = 512
filter_first_image_count = 1
filter_last_image_count = 2
filter_total_image_count = filter_first_image_count + filter_last_image_count


def read_features(file_list, feature_file, feat_dim=512, skip_offset=0):
    video_image_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 1000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/sqqjd/ycjy/150100414a54443452064d2742d03600/video/20200509_173513/00110_000.jpg
            lines = line.strip().split()
            file_path = lines[0].strip()

            feature = feat_file.read(feat_dim * 4 + skip_offset)
            feature = feature[skip_offset:]

            video_name = os.path.dirname(file_path)
            file_name = os.path.basename(file_path)
            image_features = video_image_features.get(video_name, None)
            if image_features is None:
                video_image_features[video_name] = [[file_name, feature]]
            else:
                image_features.append([file_name, feature])

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return video_image_features


def main():
    faceid_videos = {}
    # video_set = {}
    with open(face_select_list_file, 'r') as file:
        for line in file.readlines():
            # 1687 1687 gd/zs/nqjd/dxjyjy/150100414a54443452064d2742e83600/video/20190702_174057/00131_000.jpg 0
            lines = line.split()
            faceid = '{} {}'.format(lines[0].strip(), lines[1].strip())
            video = os.path.dirname(lines[2])
            videos = faceid_videos.get(faceid, None)
            if videos is None:
                faceid_videos[faceid] = {video}
            else:
                videos.add(video)
            # video_set.add(video)

    video_image_features = read_features(file_list_path, feature_file_path, feat_dim=feat_dim)

    out_path = os.path.dirname(file_list_path)

    out_list_file_path = os.path.join(out_path, 'Feature.select.txt')
    out_list_file = open(out_list_file_path, 'w')
    out_feat_file_path = os.path.join(out_path, 'Feature.select.dat')
    out_feat_file = open(out_feat_file_path, 'wb')
    out_file_path = os.path.join(out_path, 'faceid_video_list.txt')
    out_file = open(out_file_path, 'w')
    file_count = 0
    faceid_count = 0
    video_count = 0
    for faceid, videos in faceid_videos.items():
        faceid_count += 1
        videos = list(videos)
        videos.sort()
        for video in videos:
            image_features = video_image_features.get(video, None)
            if image_features is None:
                continue
            image_features.sort(key=lambda x: x[0])
            image_count = len(image_features)
            if image_count <= filter_total_image_count:
                continue
            image_features = image_features[filter_first_image_count:(image_count - filter_last_image_count)]
            for image, feature in image_features:
                out_feat_file.write(feature)
                out_info = '{} {} 0\n'.format(faceid, os.path.join(video, image))
                out_list_file.write(out_info)
                file_count += 1
                if file_count % 1000 == 0:
                    print('Proccess File Count:', file_count)
            out_info = '{} {}\n'.format(faceid, video)
            out_file.write(out_info)
            video_count += 1
    out_file.close()
    out_feat_file.close()
    out_list_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    os.system('chmod a+wr {}'.format(out_feat_file_path))
    os.system('chmod a+wr {}'.format(out_list_file_path))
    print('Total Face Count: {}, Filter Face Count: {}, Video Count: {}, File Count: {}'.format(len(faceid_videos), faceid_count, video_count, file_count))


if __name__ == '__main__':
    main()
